package com.starhub.application.rag.processor.domain;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;

import java.util.HashMap;
import java.util.Hashtable;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import com.starhub.application.collection.entity.KnowledgeCollection;
import com.starhub.application.collection.service.KnowledgeCollectionService;
import com.starhub.application.rag.constant.RagConstant;
import com.starhub.application.rag.dto.MetaRagDto;
import com.starhub.application.rag.processor.EmbeddingProcessor;
import com.starhub.application.rag.result.RetrieveResult;
import com.starhub.application.rag.service.EmbeddingService;
import com.starhub.application.rag.store.MilvusStore;
import com.starhub.application.rag.store.MilvusStore.SimilarityResult;
import com.starhub.application.rag.store.RetrieveStore;
import com.starhub.application.rag.transformer.SimpleTextSegmentTransformer;
import com.starhub.application.rag.util.MilvusConditionBuilder;
import com.starhub.integration.milvus.langchain.MilvusMetadataFilterMapper;
import com.starhub.utils.Util;
import com.starhub.application.rag.config.ThreadPoolConfig;

/**
 * 向量化处理器
 * 用于处理向量化相关的业务逻辑
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-01-27
 */
@Component
@Slf4j
public class EmbeddingMetaProcessor extends EmbeddingProcessor {

    @Autowired
    SimpleTextSegmentTransformer simpleTextSegmentTransformer;

    @Autowired
    private MilvusStore milvusStore;

    @Autowired
    private KnowledgeCollectionService knowledgeCollectionService;

    @Autowired
    private MilvusConditionBuilder milvusConditionBuilder;

    @Autowired
    private EmbeddingService embeddingService;

    @Autowired
    private ThreadPoolConfig threadPoolConfig;

    /**
     * 向量化检索 关联检索模式 返回主数据和关联度
     * 
     * @param metaRagDto
     * @param vectorModel
     * @return
     */
    public String embeddingRetrievalMeta(MetaRagDto metaRagDto, String vectorModel, Embedding questionEmbedding) {
        // 获取元数据表关联
        List<KnowledgeCollection> metaCollections = knowledgeCollectionService
                .selectBatchIds(metaRagDto.getMetaCollectionIds());

        // 执行每个元数据表的检索任务
        List<Future<RetrieveResult>> futures = metaCollections.stream()
                .map(collection -> threadPoolConfig.getThreadPoolExecutor().submit(
                        () -> {
                            return executeEmbeddingRetrieval(collection, metaRagDto, vectorModel,
                                questionEmbedding);
                        }))
                .collect(Collectors.toList());

        List<RetrieveResult> retrieveResults = futures.stream().map(future -> {
            try {
                return future.get();
            } catch (Exception e) {
                log.error("检索任务执行失败", e);
            }
            return null;
        }).collect(Collectors.toList());
        return milvusConditionBuilder.buildMetaLinkResult(retrieveResults);
    }
    

    /**
     * 执行单个集合的检索任务
     * 
     * @param collection        知识库集合
     * @param metaRagDto        元数据检索请求
     * @param vectorModel       向量模型标识
     * @param questionEmbedding 问题向量
     * @return 检索结果
     */
    private RetrieveResult executeEmbeddingRetrieval(KnowledgeCollection collection,
            MetaRagDto metaRagDto,
            String vectorModel,
            Embedding questionEmbedding) {
        try {

            // 获取检索结果
            List<String> itemIds = this.embeddingRetrieval(collection, metaRagDto, vectorModel, questionEmbedding); 

            //元数据关联模式解析
            String itemIdsStr = String.join(",", itemIds);
            //元数据id 向量化
            Embedding linkEmbedding = embeddingService.embedding(vectorModel, itemIdsStr);

            // 构建检索条件
            Filter filter = new MetadataFilterBuilder(RagConstant.LinkMetaData.META_COLLECTION)
                .isEqualTo(Util.null2String(collection.getId()));
            String condition = MilvusMetadataFilterMapper.map(filter, RagConstant.LangchainMilvusFields.FIELD_METADATA);

            Integer topK = RagConstant.RETRIEVE_EMBEDDING_COUNT;
            Double minScore = 0.0;
            String LinkCollectionName = metaRagDto.getMainCollectionId() + RagConstant.META_POSTFIX;
            
            // 根据元数据id内容 向量检索关联主数据
            List<SimilarityResult> _results = milvusStore.search(linkEmbedding, topK, minScore, condition, LinkCollectionName);

            // 获取有关联的主数据ID
            Map<String, Float> mainIdsScore = new Hashtable<>();
            List<String> mainIds = _results.stream().map(similarityResult -> {
                Metadata metadata = similarityResult.getTextSegment().metadata();
                String mainId = metadata.getString(RagConstant.LinkMetaData.MAIN_ID);
                mainIdsScore.put(mainId, similarityResult.getScore());
                return mainId;
            }).collect(Collectors.toList());

            return RetrieveResult.builder()
                        .collectionId(collection.getId())
                        .knowledgeCollection(collection)
                        .itemIds(itemIds)
                        .mainIds(mainIds)
                        .mainIdsScore(mainIdsScore)
                        .build();
        } catch (Exception e) {
            log.error("集合 {} 检索失败: {}", collection.getCollectionIdentifier(), e.getMessage(), e);
            return null;
        }
    }
}

